import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import os
import joblib

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'Arial Unicode MS']  # 尝试多种中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def load_processed_data(file_path='../data/feature_engineering/room_occupancy_with_features.csv'):
    """
    加载经过特征工程处理的数据
    
    参数:
    file_path: 处理后数据的文件路径
    
    返回:
    DataFrame: 加载的数据
    """
    print(f"加载数据: {file_path}")
    # 检查文件是否存在
    if not os.path.exists(file_path):
        print(f"文件不存在: {file_path}")
        raise FileNotFoundError(f"无法找到数据文件: {file_path}")
    
    # 加载数据
    df = pd.read_csv(file_path)
    print(f"成功加载数据，形状: {df.shape}")
    return df

def prepare_data(df):
    """
    准备训练和测试数据
    
    参数:
    df: 数据DataFrame
    
    返回:
    X_train, X_test, y_train, y_test: 训练集和测试集
    feature_names: 特征名称列表
    """
    print("\n准备训练和测试数据...")
    
    # 删除不需要的列
    time_cols = ['Date', 'Time', 'DateTime']
    time_cols = [col for col in time_cols if col in df.columns]
    if time_cols:
        df = df.drop(columns=time_cols)
        print(f"已删除时间列: {time_cols}")
    
    # 分离特征和目标变量
    X = df.drop(columns=['Room_Occupancy_Count'])
    y = df['Room_Occupancy_Count']
    
    # 保存特征名称
    feature_names = X.columns.tolist()
    print(f"特征数量: {len(feature_names)}")
    
    # 划分训练集和测试集 - 使用分层抽样
    # 对于回归问题，创建伪分层
    y_binned = pd.cut(y, bins=5, labels=False)
    
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y_binned
    )
    
    print(f"训练集大小: {X_train.shape}, 测试集大小: {X_test.shape}")
    return X_train, X_test, y_train, y_test, feature_names

def train_random_forest(X_train, y_train, tune_hyperparameters=False):
    """
    训练随机森林模型
    
    参数:
    X_train: 训练特征
    y_train: 训练标签
    tune_hyperparameters: 是否进行超参数调优
    
    返回:
    训练好的模型
    """
    print("\n开始训练随机森林模型...")
    
    if tune_hyperparameters:
        print("进行超参数调优...")
        # 定义参数网格
        param_grid = {
            'n_estimators': [100, 200, 300],
            'max_depth': [None, 15, 20, 30],
            'min_samples_split': [2, 5, 10],
            'min_samples_leaf': [1, 2, 4],
            'max_features': ['auto', 'sqrt', 'log2'],
            'bootstrap': [True, False]
        }
        
        # 创建基础模型
        rf = RandomForestRegressor(random_state=42, n_jobs=-1)
        
        # 创建网格搜索对象
        grid_search = GridSearchCV(
            estimator=rf,
            param_grid=param_grid,
            cv=5,
            n_jobs=-1,
            scoring='neg_mean_squared_error',
            verbose=1
        )
        
        # 执行网格搜索
        grid_search.fit(X_train, y_train)
        
        # 获取最佳参数和模型
        best_params = grid_search.best_params_
        model = grid_search.best_estimator_
        
        print(f"最佳参数: {best_params}")
    else:
        # 使用优化的参数
        model = RandomForestRegressor(
            n_estimators=200,
            max_depth=None,
            min_samples_split=5,
            min_samples_leaf=2,
            max_features='sqrt',
            bootstrap=True,
            random_state=42,
            n_jobs=-1  # 使用所有可用的CPU核心
        )
        
        # 训练模型
        model.fit(X_train, y_train)
    
    print("随机森林模型训练完成")
    return model

def evaluate_model(model, X_test, y_test):
    """
    评估模型性能
    
    参数:
    model: 训练好的模型
    X_test: 测试特征
    y_test: 测试标签
    
    返回:
    评估指标字典
    """
    print("\n评估随机森林模型性能...")
    
    # 预测测试集
    y_pred = model.predict(X_test)
    
    # 计算评估指标
    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    
    # 打印评估结果
    print(f"均方误差 (MSE): {mse:.4f}")
    print(f"均方根误差 (RMSE): {rmse:.4f}")
    print(f"平均绝对误差 (MAE): {mae:.4f}")
    print(f"决定系数 (R²): {r2:.4f}")
    
    # 返回评估指标
    metrics = {
        'mse': mse,
        'rmse': rmse,
        'mae': mae,
        'r2': r2
    }
    
    return metrics, y_pred

def visualize_results(y_test, y_pred, feature_names, model, output_dir='../static/model_results'):
    """
    可视化模型结果
    
    参数:
    y_test: 测试标签
    y_pred: 预测标签
    feature_names: 特征名称列表
    model: 训练好的模型
    output_dir: 输出目录
    """
    print("\n可视化模型结果...")
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 1. 预测值与实际值对比散点图
    plt.figure(figsize=(10, 6))
    plt.scatter(y_test, y_pred, alpha=0.5)
    plt.plot([0, max(y_test.max(), y_pred.max())], [0, max(y_test.max(), y_pred.max())], 'r--')
    plt.xlabel('实际值')
    plt.ylabel('预测值')
    plt.title('随机森林模型: 预测值 vs 实际值')
    plt.grid(True)
    plt.savefig(os.path.join(output_dir, '随机森林模型_预测值与实际值对比.png'))
    plt.close()
    
    # 2. 预测误差直方图
    errors = y_test - y_pred
    plt.figure(figsize=(10, 6))
    plt.hist(errors, bins=30, alpha=0.7, color='skyblue')
    plt.xlabel('预测误差')
    plt.ylabel('频率')
    plt.title('随机森林模型: 预测误差分布')
    plt.grid(True)
    plt.savefig(os.path.join(output_dir, '随机森林模型_预测误差分布.png'))
    plt.close()
    
    # 3. 特征重要性可视化
    feature_importance = model.feature_importances_
    # 创建特征重要性DataFrame
    importance_df = pd.DataFrame({
        'Feature': feature_names,
        'Importance': feature_importance
    })
    # 按重要性排序
    importance_df = importance_df.sort_values('Importance', ascending=False)
    
    # 绘制前15个最重要的特征
    plt.figure(figsize=(12, 8))
    sns.barplot(x='Importance', y='Feature', data=importance_df.head(15))
    plt.title('随机森林模型: 特征重要性 (Top 15)')
    plt.xlabel('重要性')
    plt.ylabel('特征')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, '随机森林模型_特征重要性.png'))
    plt.close()
    
    print(f"可视化结果已保存至: {output_dir}")

def save_model(model, output_dir='../model/saved_models'):
    """
    保存训练好的模型
    
    参数:
    model: 训练好的模型
    output_dir: 输出目录
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 保存模型
    model_path = os.path.join(output_dir, 'random_forest_model.joblib')
    joblib.dump(model, model_path)
    print(f"\n模型已保存至: {model_path}")

def random_forest_pipeline(tune_hyperparameters=False):
    """
    随机森林模型完整流程
    
    参数:
    tune_hyperparameters: 是否进行超参数调优
    
    返回:
    训练好的模型和评估指标
    """
    print("开始随机森林模型流程...")
    
    # 1. 加载数据
    df = load_processed_data()
    
    # 2. 准备数据
    X_train, X_test, y_train, y_test, feature_names = prepare_data(df)
    
    # 3. 训练模型
    model = train_random_forest(X_train, y_train, tune_hyperparameters)
    
    # 4. 评估模型
    metrics, y_pred = evaluate_model(model, X_test, y_test)
    
    # 5. 可视化结果
    visualize_results(y_test, y_pred, feature_names, model)
    
    # 6. 保存模型
    save_model(model)
    
    print("\n随机森林模型流程完成!")
    return model, metrics

if __name__ == "__main__":
    # 运行随机森林模型流程
    random_forest_pipeline(tune_hyperparameters=False)